import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import numpy as np
from sklearn.manifold import TSNE
from sklearn.preprocessing import StandardScaler
import argparse
import os

# 设置中文字体，以防图例乱码
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False



class FeatureVisualizer:
    def __init__(self, csv_path: str, output_dir: str = 'visualizations_v2'):
        # ... (初始化代码不变, 只是改了默认输出目录) ...
        print(f"正在加载数据: {csv_path}")
        self.df = pd.read_csv(csv_path)
        self.output_dir = output_dir
        os.makedirs(self.output_dir, exist_ok=True)
        print(f"数据加载完成: {self.df.shape[0]}行, {self.df.shape[1]}列")
        print(f"输出目录: {self.output_dir}")
        
        # 计算良恶意包总数，用于后续计算比率
        self.total_malicious = self.df[self.df['label'] == 1].shape[0]
        self.total_benign = self.df[self.df['label'] == 0].shape[0]
        print(f"  - 恶意样本数: {self.total_malicious}")
        print(f"  - 良性样本数: {self.total_benign}")


    def plot_hit_rules_distribution(self):
        """方案一：绘制规则命中总数分布图 (保留原图，因为它揭示了重要现象)"""
        print("\n[方案一] 正在生成规则命中总数分布图...")
        plt.figure(figsize=(12, 7))
        
        sns.kdeplot(data=self.df, x='hit_rules_count', hue='data_type', 
                    fill=True, common_norm=False, palette={'malicious': 'red', 'benign': 'blue'},
                    alpha=0.5, linewidth=2)
        
        plt.title('规则命中总数分布对比 (重要发现)', fontsize=16)
        plt.xlabel('单个包命中的规则总数', fontsize=12)
        plt.ylabel('密度', fontsize=12)
        plt.legend(title='数据类型', labels=['恶意 (malicious)', '良性 (benign)'])
        plt.grid(axis='y', linestyle='--', alpha=0.7)
        
        output_path = os.path.join(self.output_dir, '1_hit_rules_distribution.png')
        plt.savefig(output_path, dpi=300)
        plt.close()
        print(f"  ✅ 图表已保存: {output_path}")
        print("  💡 观察: 此图揭示了良性包可能比恶意包更复杂，触发更多通用规则的现象。")

    def plot_rule_discriminative_power(self):
        """优化方案 2.1：绘制规则区分能力散点图"""
        print("\n[优化方案 2.1] 正在生成规则区分能力散点图...")
        rule_cols = [col for col in self.df.columns if col.startswith('rule_')]
        
        rule_summary = []
        for col in rule_cols:
            malicious_hits = self.df[(self.df[col] == 1) & (self.df['label'] == 1)].shape[0]
            benign_hits = self.df[(self.df[col] == 1) & (self.df['label'] == 0)].shape[0]
            
            mal_rate = malicious_hits / self.total_malicious if self.total_malicious > 0 else 0
            ben_rate = benign_hits / self.total_benign if self.total_benign > 0 else 0
            
            if mal_rate + ben_rate > 0: # 只显示至少命中过一次的规则
                rule_summary.append({
                    'rule_id': col.replace('rule_', ''),
                    'malicious_rate': mal_rate * 100,
                    'benign_rate': ben_rate * 100
                })
        
        summary_df = pd.DataFrame(rule_summary)
        
        plt.figure(figsize=(12, 12))
        g = sns.scatterplot(data=summary_df, x='benign_rate', y='malicious_rate', 
                            s=100, alpha=0.7, color='purple')
        
        # 添加辅助线和区域说明
        plt.axline((0, 0), slope=1, color='gray', linestyle='--', label='区分度=1')
        plt.fill_between([0, 100], [0, 100], 100, color='red', alpha=0.1, label='恶意倾向区')
        plt.fill_between([0, 100], 0, [0, 100], color='blue', alpha=0.1, label='良性倾向区')

        # 标注一些关键规则
        for i, point in summary_df.iterrows():
            if point['malicious_rate'] > 50 or point['benign_rate'] > 50 or (point['malicious_rate'] > 20 and point['benign_rate'] < 5):
                g.text(point['benign_rate']+0.5, point['malicious_rate'], point['rule_id'], fontsize=9)

        plt.title('规则区分能力散点图 (Rule Discriminative Power)', fontsize=16)
        plt.xlabel('良性命中率 (%)', fontsize=12)
        plt.ylabel('恶意命中率 (%)', fontsize=12)
        plt.xlim(-2, 102)
        plt.ylim(-2, 102)
        plt.grid(True, linestyle='--', alpha=0.6)
        plt.legend()
        
        output_path = os.path.join(self.output_dir, '2_rule_discriminative_power.png')
        plt.savefig(output_path, dpi=300)
        plt.close()
        print(f"  ✅ 图表已保存: {output_path}")

    def plot_graph_feature_power(self):
        """优化方案 3.1：绘制图特征区分能力条形图"""
        print("\n[优化方案 3.1] 正在生成图特征区分能力条形图...")
        
        try:
            from feature_extractor import _get_empty_graph_features
            graph_feature_names = list(_get_empty_graph_features().keys())
        except ImportError:
            print("  ⚠️ 无法从 feature_extractor 导入图特征列表，将从数据中推断。")
            graph_feature_names = [c for c in self.df.columns if not c.startswith('rule_') and c not in ['package_name', 'label', 'data_type', 'hit_rules_count']]

        graph_feature_names = [f for f in graph_feature_names if f in self.df.columns]

        if not graph_feature_names:
            print("  ⚠️ 未找到图特征，跳过此图。")
            return

        power_scores = []
        epsilon = 1e-9 # 防止除以零
        
        for feature in graph_feature_names:
            mean_malicious = self.df[self.df['label'] == 1][feature].mean()
            mean_benign = self.df[self.df['label'] == 0][feature].mean()
            
            score = np.log2((mean_malicious + epsilon) / (mean_benign + epsilon))
            power_scores.append({'feature': feature, 'score': score})
            
        power_df = pd.DataFrame(power_scores).sort_values('score', ascending=False)
        
        plt.figure(figsize=(12, 8))
        colors = ['red' if x > 0 else 'blue' for x in power_df['score']]
        sns.barplot(data=power_df, x='score', y='feature', palette=colors)
        
        plt.title('图特征区分能力 (Discriminative Power)', fontsize=16)
        plt.xlabel('区分度得分 (log2(恶意均值 / 良性均值))', fontsize=12)
        plt.ylabel('图特征', fontsize=12)
        plt.axvline(0, color='black', linewidth=0.8)
        plt.grid(axis='x', linestyle='--', alpha=0.7)
        
        output_path = os.path.join(self.output_dir, '3_graph_feature_power.png')
        plt.savefig(output_path, dpi=300)
        plt.close()
        print(f"  ✅ 图表已保存: {output_path}")

    def plot_tsne_clustering(self, feature_type='all'):
        """方案四：绘制特征空间降维聚类图 (保持不变，因为效果很好)"""
        # ... (此函数代码与之前版本完全相同) ...
        print(f"\n[方案四] 正在生成t-SNE降维聚类图 (特征类型: {feature_type})...")
        
        if feature_type == 'all':
            try:
                from feature_extractor import _get_empty_graph_features
                graph_feature_names = list(_get_empty_graph_features().keys())
            except ImportError:
                graph_feature_names = [c for c in self.df.columns if not c.startswith('rule_') and c not in ['package_name', 'label', 'data_type', 'hit_rules_count']]
            feature_cols = [col for col in self.df.columns if col.startswith('rule_') or col in graph_feature_names]
        elif feature_type == 'rules':
            feature_cols = [col for col in self.df.columns if col.startswith('rule_')]
        elif feature_type == 'graph':
            try:
                from feature_extractor import _get_empty_graph_features
                graph_feature_names = list(_get_empty_graph_features().keys())
            except ImportError:
                graph_feature_names = [c for c in self.df.columns if not c.startswith('rule_') and c not in ['package_name', 'label', 'data_type', 'hit_rules_count']]
            feature_cols = [f for f in graph_feature_names if f in self.df.columns]
        else:
            print(f"  ⚠️ 未知特征类型: {feature_type}")
            return
            
        if not feature_cols:
            print(f"  ⚠️ 未找到 '{feature_type}' 类型的特征，跳过。")
            return

        features = self.df[feature_cols].fillna(0)
        labels = self.df['data_type']
        
        # 标准化
        scaler = StandardScaler()
        features_scaled = scaler.fit_transform(features)
        
        print("  正在进行t-SNE计算 (这可能需要一些时间)...")
        tsne = TSNE(n_components=2, random_state=42, perplexity=min(30, len(features)-1), n_iter=1000)
        tsne_results = tsne.fit_transform(features_scaled)
        
        tsne_df = pd.DataFrame(data=tsne_results, columns=['tsne-2d-one', 'tsne-2d-two'])
        tsne_df['data_type'] = labels
        
        plt.figure(figsize=(12, 10))
        sns.scatterplot(
            x="tsne-2d-one", y="tsne-2d-two",
            hue="data_type",
            palette={'malicious': 'red', 'benign': 'blue'},
            data=tsne_df,
            legend="full",
            alpha=0.6
        )
        
        plt.title(f'特征空间t-SNE降维可视化 ({feature_type.capitalize()} Features)', fontsize=16)
        plt.xlabel('t-SNE维度1', fontsize=12)
        plt.ylabel('t-SNE维度2', fontsize=12)
        
        output_path = os.path.join(self.output_dir, f'4_tsne_clustering_{feature_type}.png')
        plt.savefig(output_path, dpi=300)
        plt.close()
        print(f"  ✅ 图表已保存: {output_path}")

    def run_all(self):
        """运行所有可视化方案"""
        self.plot_hit_rules_distribution()
        self.plot_rule_discriminative_power()
        self.plot_graph_feature_power()
        self.plot_tsne_clustering(feature_type='all')
        # self.plot_tsne_clustering(feature_type='rules') # 可以根据需要取消注释
        # self.plot_tsne_clustering(feature_type='graph') # 可以根据需要取消注释
        print("\n🎉 所有优化后的可视化图表生成完毕！")

def main():
    parser = argparse.ArgumentParser(description="特征集可视化分析工具 (v2 - 优化版)")
    parser.add_argument('csv_file', help='包含特征和标签的CSV文件路径')
    parser.add_argument('--output', '-o', default='visualizations_v2', help='可视化图片输出目录')
    args = parser.parse_args()
    
    visualizer = FeatureVisualizer(args.csv_file, args.output)
    visualizer.run_all()

if __name__ == '__main__':
    main()